#include "CodeGenerator.h"
#include "neiderra/utils/Lexer.h"
#include "neiderra/core/io/IOStream.h"
#include "VirtualMachine.h"
#include "neiderra/utils/CommonFunctions.h"

#include <string>
#include <vector>
#include <algorithm>
#include <map>

#define log(s) core->getLogger()->log(s)
#define trace(s) core->getLogger()->traceBlock(s)
#define leave core->getLogger()->leaveBlock()

using namespace neiderra::vm;
using namespace neiderra::types;
using namespace neiderra::utils::lexer;
using namespace neiderra::utils;
using std::string;
using std::map;
using std::vector;

CodeGenerator::CodeGenerator(Core* core)
	: Base(core)
{
	_lexer = new Lexer(core);
	_lexer->skipComments(csASM);
}

CodeGenerator::~CodeGenerator()
{
	delete _lexer;
}

CodeGenerator::Command CodeGenerator::getCommand(const string& cmd)
{
	if(cmd == "add")
	{
		return CodeGenerator::cmdADD;
	}
	else if(cmd == "sub")	
	{
		return CodeGenerator::cmdSUB;
	}
	else if(cmd == "div")
	{
		return CodeGenerator::cmdDIV;
	}
	else if(cmd == "mul")
	{
		return CodeGenerator::cmdMUL;
	}
	else if(cmd == "nop")
	{
		return CodeGenerator::cmdNOP;
	}
	else if(cmd == "out")
	{
		return CodeGenerator::cmdOUT;
	}
	else if(cmd == "st")
	{
		return CodeGenerator::cmdST;
	}
	else if(cmd == "mov")
	{
		return CodeGenerator::cmdMOV;
	}
	else if(cmd == "push")
	{
		return CodeGenerator::cmdPUSH;
	}
	else if(cmd == "pop")
	{
		return CodeGenerator::cmdPOP;
	}
	else if(cmd == "top")
	{
		return CodeGenerator::cmdTOP;
	}
	else if(cmd == "ld")
	{
		return CodeGenerator::cmdLD;
	}
	else if(cmd == "goto")
	{
		return CodeGenerator::cmdGOTO;
	}
	else if(cmd == "db")
	{
		return CodeGenerator::cmdDB;
	}
	else if(cmd == "org")
	{
		return CodeGenerator::cmdORG;
	}
	else if(cmd == "breq")
	{
		return CodeGenerator::cmdBREQ;
	}
	else if(cmd == "brneq")
	{
		return CodeGenerator::cmdBRNEQ;
	}
	else if(cmd == "inc")
	{
		return CodeGenerator::cmdINC;
	}
	else if(cmd == "dec")
	{
		return CodeGenerator::cmdDEC;
	}
	else if(cmd == "cp")
	{
		return CodeGenerator::cmdCP;
	}
	else if(cmd == "call")
	{
		return CodeGenerator::cmdCALL;
	}
	else if(cmd == "proc")
	{
		return CodeGenerator::cmdPROC;
	}
	else if(cmd == "ret")
	{
		return CodeGenerator::cmdRET;
	}
	else if(cmd == "pushm")
	{
		return CodeGenerator::cmdPUSHM;
	}
	else if(cmd == "popm")
	{
		return CodeGenerator::cmdPOPM;
	}
	else if(cmd == "topm")
	{
		return CodeGenerator::cmdTOPM;
	}
	else if(cmd == "dbm")
	{
		return CodeGenerator::cmdDBM;
	}
	else if(cmd == "stm")
	{
		return CodeGenerator::cmdSTM;
	}
	
	else return CodeGenerator::cmdERR;
}

void CodeGenerator::match(tokens t)
{
	if(_ct.id != t)
		throw exBadToken(_lexer->getLine());
	_ct = _lexer->nextToken();
}

bool CodeGenerator::next(tokens tok)
{
	return (_ct.id==tok)?true:false;
}

unsigned CodeGenerator::getRegNum(const string& regName)
{
	if(regName[0]!='r')
		throw exBadRegister();
	char buf[3]={0,0,0};
	regName.copy(buf, 3, 1);					
	unsigned regnum = (unsigned)atoi((const char*)buf);
    if(regnum > 32)
		throw exBadRegister();
	return regnum;
}

string CodeGenerator::extractParentName(const string& name)
{
	char buf[80];
	size_t bobosInd = name.find_first_of('$', 1);
	name.copy(buf, bobosInd, 0);
	buf[bobosInd]='\0';
	return buf;
}

unsigned CodeGenerator::getParentOffset(const string& name)
{
	string parent_name = extractParentName(name);
	return _symtab[parent_name];
}

unsigned CodeGenerator::getAbsoluteOffset(const string& name)
{
	const char* p = name.c_str();
	return count(p, p + name.size(), '$') - 1;
}

void CodeGenerator::writeInt(int n)
{
	nByte* num = reinterpret_cast<nByte*>(&n);
	writeBytes(num, 4);
}

void CodeGenerator::writeUInt(unsigned int n)
{
	nByte* num = reinterpret_cast<nByte*>(&n);
	writeBytes(num, 4);
}

void CodeGenerator::writeByte(nByte b)
{
//	Core::defLogger->log((int)b);
	_bytecode[_counter] = b;
	_counter ++;
}

void CodeGenerator::writeBytes(nByte* bytes, unsigned size)
{
	for(int i=0;i<size;i++,_counter++){
		nByte val = bytes[i];
//		Core::defLogger->log((int)val);
		_bytecode[_counter] = bytes[i];
	}
}

void CodeGenerator::getId(string& id_name)
{
	// identifier begins with '$'
	tokens lasttok = tokID;
	
	do{
		if(_ct.id == tokID){
			if(lasttok == tokID) 
				break;
            id_name.append(_ct.value.sval);
			lasttok = tokID;
		}
		else if(_ct.id == tokBOBOS)
		{
			if(lasttok == tokBOBOS)
				break;
			id_name.append("$");
			lasttok = tokBOBOS;
		}
		else break;
		_ct = _lexer->nextToken();
	}while(1);
}

nByte* CodeGenerator::generate(iInputStream* input, unsigned& size)
{
	trace("CodeGenerator::generate");
	nByte* bytecode = new nByte[4096];
	_bytecode = bytecode;
	
	
	// place for org
	neiderra::utils::memset(_bytecode, 0, 5);
	_counter = 5;
	_mainpos = 0;

	_input = input;
	_lexer->setInputStream(input);

	_ct = _lexer->nextToken();
	
	try {
		_procs["org"] = 0x00;
		while(!next(tokEOFTOK))
		{
			// get cmd 
			string scmd = _ct.value.sval;
			log(scmd);
				
			Command cmd = getCommand(scmd);

			match(tokID);
			
			// look for a label
			if(next(tokDOT_DOT) && cmd == cmdERR)
			{
				_procs[scmd]=_counter;
//				_counter++;
				match(tokDOT_DOT);
				continue;
			}
			
			if (cmd == cmdERR)
			    throw exBadCommand();

			if(cmd == cmdADD)
			{
				nByte size = (nByte)_ct.value.num_int;
				match(tokNUM_INTEGER);
				
				writeByte(VirtualMachine::cmdADD);
				writeByte(size);
			}
			else if(cmd == cmdMUL)
			{
				nByte size = (nByte)_ct.value.num_int;
				match(tokNUM_INTEGER);
				
				writeByte(VirtualMachine::cmdMUL);
				writeByte(size);
			}
			else if(cmd == cmdDIV)
			{
				nByte size = (nByte)_ct.value.num_int;
				match(tokNUM_INTEGER);
				
				writeByte(VirtualMachine::cmdDIV);
				writeByte(size);
			}
			else if(cmd == cmdSUB)
			{
				nByte size = (nByte)_ct.value.num_int;
				match(tokNUM_INTEGER);
				
				writeByte(VirtualMachine::cmdSUB);
				writeByte(size);
			}
			else if(cmd == cmdDEC || cmd == cmdINC)
			{
				nByte reg = (nByte) getRegNum(_ct.value.sval);
				match(tokID);
				if(reg > 32)
                    throw exBadRegister();

				if(cmd == cmdINC)
					writeByte(VirtualMachine::cmdINC);
				else if( cmd == cmdDEC )
					writeByte(VirtualMachine::cmdDEC);

				writeByte(reg);
			}
			else if(cmd == cmdDB)
			{
				string id_name;
				getId(id_name);

				match(tokCOMMA);
				nByte bytes_num = _ct.value.num_int;
				match(tokNUM_INTEGER);

				string parent_name = extractParentName(id_name);

				// id->offset = f$id->offset
				// f->offset += sizeof(id)
				unsigned parent_offset = _symtab[parent_name];
				_symtab[id_name] = parent_offset;
				_symtab[parent_name] += bytes_num;

				writeByte(VirtualMachine::cmdDATA);
                writeByte(bytes_num);
			}
			else if(cmd == cmdST)
			{
				unsigned offset;
				if(next(tokNUM_INTEGER))	// tokNUM_HEX
				{
                    offset = _ct.value.num_int;	// !
				}
				else if(next(tokID) || next(tokBOBOS))	// variable
				{
					// if variable then get offset from symtable
					string id_name;
					getId(id_name);
                    offset = _symtab[id_name];
				}

				match(tokCOMMA);
				nByte reg_num = (nByte) getRegNum(_ct.value.sval);
                match(tokID);				
				match(tokCOMMA);
				nByte size = (nByte) _ct.value.num_int;
				match(tokNUM_INTEGER);

				writeByte(VirtualMachine::cmdST);
				// relative offset
				writeUInt(offset);
				writeByte(reg_num);
				writeByte(size);
			}
			else if(cmd == cmdCP || cmd == cmdMOV)
			{
				nByte reg1 = (nByte) getRegNum(_ct.value.sval);
				match(tokID);
				match(tokCOMMA);
				nByte reg2 = (nByte) getRegNum(_ct.value.sval);
				match(tokID);

				if(reg1 > 32 || reg2 > 32 )
					throw exBadRegister();

				if(cmd == cmdCP)
					writeByte(VirtualMachine::cmdCP);
				else writeByte(VirtualMachine::cmdMOV);
				writeByte(reg1);
				writeByte(reg2);
			}
			else if(cmd == cmdNOP)
			{
				writeByte(VirtualMachine::cmdNOP);
			}
			else if(cmd == cmdLD)
			{
				unsigned regNum = getRegNum(_ct.value.sval);
				//nByte reg = (nByte)_ct.value.num_int;
				//match(tokNUM_INTEGER);
				match(tokID);
				match(tokCOMMA);
				int val = _ct.value.num_int;
				match(tokNUM_INTEGER);
				match(tokCOMMA);
				nByte size = (nByte)_ct.value.num_int;
				match(tokNUM_INTEGER);

				writeByte(VirtualMachine::cmdLD);
				writeByte((nByte)regNum);
				writeInt(val);
				writeByte(size);
			}
			else if(cmd == cmdPROC)
			{
				string proc_name;
				getId(proc_name);
				_procs[proc_name]=_counter;		// for goto
				if(proc_name == "$main")
					_mainpos = _counter;		// for entry point
				
				match(tokDOT_DOT);
			}
			else if(cmd == cmdORG)
			{
				unsigned addr = _ct.value.num_int;
				writeByte(VirtualMachine::cmdGOTO);
				writeUInt(addr);

				match(tokNUM_INTEGER);	// tokNUM_HEX
			}
			else if(cmd == cmdCALL || cmd == cmdBREQ
				|| cmd == cmdGOTO || cmd == cmdBRNEQ)
			{
				if(cmd == cmdCALL)
					writeByte(VirtualMachine::cmdCALL);
				else if(cmd == cmdBREQ)
					writeByte(VirtualMachine::cmdBREQ);
				else if(cmd == cmdGOTO)
					writeByte(VirtualMachine::cmdGOTO);
				else if(cmd == cmdBRNEQ)
					writeByte(VirtualMachine::cmdBRNEQ);

/*				if(!next(tokID) || !next(tokBOBOS))
					throw exBadToken(_lexer->getLine());*/
				
				string procName;// = _ct.value.sval;
				getId(procName);
				unsigned addr =_procs[procName];
				_procLabels[procName].push_back(_counter);
				writeUInt(addr);
                //match(tokID);					
			}
			else if(cmd == cmdRET)
			{
				writeByte(VirtualMachine::cmdRET);
			}
			else if(cmd == cmdPUSH || cmd == cmdPOP || cmd == cmdTOP || 
							 cmd == cmdPUSHM || cmd == cmdTOPM || cmd == cmdPOPM)
			{	
				int num_int;
				float num_float;
				int offset=0;
				if(next(tokID))	// если кладем регистр
				{
                    string regname = _ct.value.sval;
					nByte regnum = (nByte)getRegNum(regname);
					match(tokID);
					if(cmd == cmdPUSH)
						writeByte(VirtualMachine::cmdPUSHR);
					else if(cmd == cmdPOP)
						writeByte(VirtualMachine::cmdPOPR);
					else if(cmd == cmdTOP)
						writeByte(VirtualMachine::cmdTOPR);
//					writeByte(cmd == cmdPUSH? VirtualMachine::cmdPUSHR: VirtualMachine::cmdPOPR);
					writeByte(regnum);
					match(tokCOMMA);
					nByte size = (nByte)_ct.value.num_int;
					writeByte(size);
					match(tokNUM_INTEGER);
					continue;
				}
				else if(next(tokBOBOS))	// если запихиваем зн-е переменной
				{
					string id;
					getId(id);
					
					// push
					if(cmd == cmdPUSH)
						writeByte(VirtualMachine::cmdPUSH);
					else if(cmd == cmdPOP)
						writeByte(VirtualMachine::cmdPOP);
					else if(cmd == cmdTOP)
						writeByte(VirtualMachine::cmdTOP);
					else if(cmd == cmdTOPM){
						writeByte(VirtualMachine::cmdTOPM);
					}
					else if(cmd == cmdPUSHM){
						writeByte(VirtualMachine::cmdPUSHM);
					}
					else if(cmd == cmdPOPM){
						writeByte(VirtualMachine::cmdPOPM);
					}
					
					match(tokCOMMA);
					
					offset = _symtab[id];
					writeUInt(offset);	
					
					// var size
					nByte size = (nByte)_ct.value.num_int;	
					writeByte(size);
					
					match(tokNUM_INTEGER);
					continue;
				}
/*				else if(next(tokNUM_INTEGER))
				{
					numsize = 4;
					num_int = _ct.value.num_int;
					number = reinterpret_cast<nByte*>(&(num_int));
					match(tokNUM_INTEGER);
				}
				else if(next(tokNUM_FLOAT))
				{
					num_float = _ct.value.num_float;
					numsize = sizeof(float);
					number = reinterpret_cast<nByte*>(&(num_float));
					match(tokNUM_FLOAT);
				}*/
/*				match(tokCOMMA);
				nByte size = (nByte)_ct.value.num_int;
				
				writeUInt(_symtab)			// вообще-то сначала size
				writeByte(size);

				match(tokNUM_INTEGER);*/
			}
			else if(cmd == cmdOUT)
			{
				string id_name;
				getId(id_name);
				match(tokCOMMA);
				nByte size = (nByte) _ct.value.num_int;
				match(tokNUM_INTEGER);
			
				writeByte(VirtualMachine::cmdLOG);
				writeUInt(_symtab[id_name]);
				writeByte(size);
			}
			else 
				throw exNotImplemented();

		}

		// set program size
		size = _counter+1;

		map<string, unsigned>::iterator i = _procs.begin();
		for(;i!=_procs.end();i++)
		{

			const string &proc_name = i->first;
			unsigned label_addr = i->second;
			vector<unsigned>::iterator j = (_procLabels[proc_name]).begin();
			vector<unsigned>::iterator j_end = (_procLabels[proc_name]).end();
			
			for(;j!=j_end;j++)
			{
				_counter = *j;
                writeUInt(label_addr);
			}
		}
		
		//entry point
		if(_mainpos != 0) 
		{
			_counter = 0;
			writeByte(VirtualMachine::cmdGOTO);
			writeUInt(_mainpos);
		}
	}
	catch(exBadToken e)
	{
		string err = "bad token, line ";
		err += intToStr(e.line);
		core->getLogger()->fatal(err.c_str());
		leave;
		return 0;
	}
	catch(exBadCommand e)
	{
	}
	catch(exNotImplemented e)
	{
	}
	catch(exBadOpSize e)
	{
	}

	leave;
	return _bytecode;
}

/*
	команды, требующие полной адресации (точка отсчета + смещение):
	push, pop, st
	
	точка отсчета - это глубина вложенности активации, посчитанная **начиная с начала стека**, 
а не с конца. это устраняет проблему отсчета, когда неск. раз выполняется рекурсивная процедура. 
адрес точки отсчета в этом случае будет указывать на некий адрес из последней рекурсивной активации.

	class A {
		int x;				// A#x
		int f(int x)		// A#f$x
		{
			char a;			// A#f$a
		}
	};
	
	int g()
	{
		A a;
		a.x = 2;
		a.f(4);
	}
*/
